import torch
from .base_model import BaseModel
from . import networks


class Pix2PixModel(BaseModel):
    """ This class implements the pix2pix model, for learning a mapping from input images to output images given paired data.
        这个类用于实现 pix2pix 模型，来学习从图像到图像的一对数据
    The model training requires '--dataset_mode aligned' dataset.
    这个模型的训练需要 '--dataset_mode aligned' 数据集
    By default, it uses a '--netG unet256' U-Net generator,
    默认情况下，它使用了一个 '--netG unet256' U-Net 生成器
    a '--netD basic' discriminator (PatchGAN),
    一个 '--netD basic' 鉴别器 (PatchGAN)。
    and a '--gan_mode' vanilla GAN loss (the cross-entropy objective used in the orignal GAN paper).
    和一个 '--gan_mode' vanilla GAN 损失 （GAN 原论文中用到的交叉熵目标函数）

    pix2pix paper: https://arxiv.org/pdf/1611.07004.pdf
    """

    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.
            增加新的模型专用选项，并重写已存在选项的默认值
        Parameters:
            parser          -- original option parser
                                原选项分析器
            is_train (bool) -- whether training phase or test phase.
                                You can use this flag to add training-specific or test-specific options.
                                是训练解析器还是测试解析器，你可以用这个标志添加训练专用或者测试专用选项

        Returns:
            the modified parser.
            修改过得处理器

        For pix2pix, we do not use image buffer
        在 pix2pix中，我们不使用 image buffer （图像缓存）
        The training objective is: GAN Loss + lambda_L1 * ||G(A)-B||_1
        训练目标是：GAN Loss + lambda_L1 * ||G(A)-B||_1
        By default, we use vanilla GAN loss, UNet with batchnorm, and aligned datasets.
        默认情况下，我们使用 vanilla GAN 损失，有 batchnorm 的 UNet，以及配对过的数据集
        """
        # changing the default values to match the pix2pix paper (https://phillipi.github.io/pix2pix/)
        # 为匹配 pix2pix 论文(https://phillipi.github.io/pix2pix/)，改变默认值
        parser.set_defaults(norm="batch", netG="unet_256", dataset_mode="aligned")
        if is_train:
            parser.set_defaults(pool_size=0, gan_mode="vanilla")
            parser.add_argument(
                "--lambda_L1",
                type=float,
                default=100.0,
                help="weight for L1 loss L1 损失的权重",
            )

        return parser

    def __init__(self, opt):
        """Initialize the pix2pix class.
            初始化 pix2pix 类
        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
                                    存储所有的实验标志需要是BaseOptions的子类
        """
        BaseModel.__init__(self, opt)
        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        # 指定你希望打印出的训练损失。training/test 脚本会调用 <BaseModel.get_current_losses>
        self.loss_names = ["G_GAN", "G_L1", "D_real", "D_fake"]
        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
        # 指定你想存储/显示的图片。training/test 脚本会调用 <BaseModel.get_current_visuals>
        self.visual_names = ["real_A", "fake_B", "real_B"]
        # specify the models you want to save to the disk.
        # 指定你希望存储在硬盘上的模型
        # The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
        # training/test 脚本会调用 <BaseModel.save_networks> 和 <BaseModel.load_networks>
        if self.isTrain:
            self.model_names = ["G", "D"]
        else:  # during test time, only load G 在测试时，只加载 G
            self.model_names = ["G"]
        # define networks (both generator and discriminator)
        # 定义网络（生成器 和 鉴别器 公用）
        self.netG = networks.define_G(
            opt.input_nc,
            opt.output_nc,
            opt.ngf,
            opt.netG,
            opt.norm,
            not opt.no_dropout,
            opt.init_type,
            opt.init_gain,
            self.gpu_ids,
        )
        # define a discriminator;
        # 定义鉴别器；
        # conditional GANs need to take both input and output images;
        # 条件GAN同时需要获取输入和输出图像；
        # Therefore, #channels for D is input_nc + output_nc
        # 因此，D的通道数是 input_nc + output_nc （输入通道 + 输出通道）
        if self.isTrain:
            self.netD = networks.define_D(
                opt.input_nc + opt.output_nc,
                opt.ndf,
                opt.netD,
                opt.n_layers_D,
                opt.norm,
                opt.init_type,
                opt.init_gain,
                self.gpu_ids,
            )

        if self.isTrain:
            # define loss functions
            # 定义损失函数
            self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
            self.criterionL1 = torch.nn.L1Loss()
            # initialize optimizers;
            # 初始化优化器
            # schedulers will be automatically created by function <BaseModel.setup>.
            # 编排器会被函数 <BaseModel.setup> 自动创建
            self.optimizer_G = torch.optim.Adam(
                self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)
            )
            self.optimizer_D = torch.optim.Adam(
                self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)
            )
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)

    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.
            将输入数据解包并部署必要的预处理步骤
        Parameters:
            input (dict): include the data itself and its metadata information.
                            包含数据迭代器和元数据信息

        The option 'direction' can be used to swap images in domain A and domain B.
        'direction' 选项可以被用来交换 A域 和 B域 的图片
        """
        AtoB = self.opt.direction == "AtoB"
        self.real_A = input["A" if AtoB else "B"].to(self.device)
        self.real_B = input["B" if AtoB else "A"].to(self.device)
        self.image_paths = input["A_paths" if AtoB else "B_paths"]

    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>.
            运行向前步骤，被 <optimize_parameters> 和 <test> 两个函数调用
        """
        self.fake_B = self.netG(self.real_A)  # G(A)

    def backward_D(self):
        """Calculate GAN loss for the discriminator
            计算鉴别器的 GAN 损失
        """
        # Fake; stop backprop to the generator by detaching fake_B
        # 假； 通过分离 fake_B（假B）来终止反向传播到生成器
        # we use conditional GANs; we need to feed both input and output to the discriminator
        # 我们使用条件GAN；我们需要将输入和输出都喂给鉴别器
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)

        pred_fake = self.netD(fake_AB.detach())
        self.loss_D_fake = self.criterionGAN(pred_fake, False)
        # Real
        # 真
        real_AB = torch.cat((self.real_A, self.real_B), 1)
        pred_real = self.netD(real_AB)
        self.loss_D_real = self.criterionGAN(pred_real, True)
        # combine loss and calculate gradients
        # 结合损失并计算梯度
        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
        self.loss_D.backward()

    def backward_G(self):
        """Calculate GAN and L1 loss for the generator
            计算生成器的 GAN 和 L1 损失
        """
        # First, G(A) should fake the discriminator
        # 首先，G(A) 需要欺骗鉴别器
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        pred_fake = self.netD(fake_AB)
        self.loss_G_GAN = self.criterionGAN(pred_fake, True)
        # Second, G(A) = B
        # 然后， 令 G(A) = B
        self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1
        # combine loss and calculate gradients
        # 结合损失并计算梯度
        self.loss_G = self.loss_G_GAN + self.loss_G_L1
        self.loss_G.backward()

    def optimize_parameters(self):
        # compute fake images: G(A)
        # 计算假图像 G(A)
        self.forward()
        # update D
        # 更新 D

        # enable backprop for D
        # 使能 D 的反向传播
        self.set_requires_grad(self.netD, True)
        # set D's gradients to zero
        # 设置 D 的梯度为零
        self.optimizer_D.zero_grad()
        # calculate gradients for D
        # 计算 D 的梯度
        self.backward_D()
        # update D's weights
        # 更新 D 的权重
        self.optimizer_D.step()

        # update G
        # 更新 G
        
        # D requires no gradients when optimizing G
        # 当优化 G 时 D 要求不能有梯度
        self.set_requires_grad(self.netD, False)
        # set G's gradients to zero
        # 设置 G 的梯度为零
        self.optimizer_G.zero_grad()
        # calculate graidents for G
        # 计算 G 的梯度
        self.backward_G()
        # udpate G's weights
        # 更新 G 的梯度
        self.optimizer_G.step()
